#!/usr/bin/python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import re
import json
import functools
import atomicwrites
from .util import Util, Expression
from .errors import ConfigError
from .const import MachineType, OfflineReason


class Handy:

    @staticmethod
    def doNothing(msg):
        pass

    @staticmethod
    def componentGetDeviceByType(component, deviceType):
        ret = None
        for dev in component.get_devices():
            if dev.get_type() == deviceType:
                assert ret is None
                ret = dev
        assert ret is not None
        return ret

    @staticmethod
    def findItemsByComponent(items, component):
        retList = []

        for item in items:
            vendor = item.get("vendor", None)
            model = item.get("model", None)
            component_type = item.get("component_type", None)
            component_id = item.get("id", None)

            if vendor is not None:
                assert model is not None
                assert component_type is None
                keyFields = ["vendor", "model"]
            else:
                assert model is None
                assert component_type is not None
                keyFields = ["component_type"]

            if component_id is not None:
                keyFields.append("id")

            bFound = True
            for key in keyFields:
                if eval("component.get_%s()" % (key)) != eval("%s" % (key)):
                    bFound = False
            if bFound:
                retList.append(item)

        return retList

    @staticmethod
    def findComponentsByItem(components, item):
        vendor = item.get("vendor", None)
        model = item.get("model", None)
        component_type = item.get("component_type", None)
        component_id = item.get("id", None)

        if vendor is not None:
            assert model is not None
            assert component_type is None
            keyFields = ["vendor", "model"]
        else:
            assert model is None
            assert component_type is not None
            keyFields = ["component_type"]

        if component_id is not None:
            keyFields.append("id")

        objList = []
        for c in components:
            bGood = True
            for key in keyFields:
                if eval("c.get_%s()" % (key)) != eval("%s" % (key)):
                    bGood = False
                    break
            if bGood:
                objList.append(c)

        return objList

    @staticmethod
    def findItemsByDevice(items, device):
        retList = []

        for item in items:
            vendor = item.get("vendor", None)
            model = item.get("model", None)
            device_type = item.get("device_type", None)
            device_id = item.get("id", None)

            if vendor is not None:
                assert model is not None
                assert device_type is None
                keyFields = ["vendor", "model"]
            else:
                assert model is None
                assert device_type is not None
                keyFields = ["device_type"]

            if device_id is not None:
                keyFields.append("id")

            bFound = True
            for key in keyFields:
                if eval("device.get_%s()" % (key)) != eval("%s" % (key)):
                    bFound = False
            if bFound:
                retList.append(item)

        return retList

    @staticmethod
    def findDevicesByItem(devices, item):
        vendor = item.get("vendor", None)
        model = item.get("model", None)
        device_type = item.get("device_type", None)
        device_id = item.get("id", None)

        if vendor is not None:
            assert model is not None
            assert device_type is None
            keyFields = ["vendor", "model"]
        else:
            assert model is None
            assert device_type is not None
            keyFields = ["device_type"]

        if device_id is not None:
            keyFields.append("id")

        objList = []
        for d in devices:
            bGood = True
            for key in keyFields:
                if eval("d.get_%s()" % (key)) != eval("%s" % (key)):
                    bGood = False
                    break
            if bGood:
                objList.append(d)

        return objList


class MachineInfoFile:

    PATH = "/etc/machine-info"

    # cache for getMachineInfo()
    _machineInfo = None

    @staticmethod
    def checkMachineName(machine_name):
        if machine_name is None:
            return False
        return True

    @classmethod
    def getMachineInfo(cls):
        if cls._machineInfo is not None:
            return cls._machineInfo

        ret = dict()
        if os.path.exists(cls.PATH):
            with open(cls.PATH, "r") as f:
                for line in f.read().split("\n"):
                    if line.startswith("#"):
                        continue
                    m = re.fullmatch("(.*?)=(.*)", line)
                    if m is None:
                        continue
                    ret[m.group(1)] = m.group(2).strip("\"")

        cls._machineInfo = ret
        return ret

    @classmethod
    def getChassisType(cls):
        value = cls.getMachineInfo().get("CHASSIS", None)
        if value is None:
            return None

        patternDict = {
            "desjtop": MachineType.DESKTOP,       # keys are regex patterns
            "laptop": MachineType.LAPTOP,
            "tablet": MachineType.TABLET,
            "convertible": MachineType.CONVERTIBLE,
            "handset": MachineType.HANDSET,
            "watch": MachineType.WATCH,
            "server": MachineType.SERVER,
            "vm": MachineType.VM,
            "container": MachineType.CONTAINER,
        }
        for pattern, chassisType in patternDict.items():
            if re.fullmatch(pattern, value):
                return chassisType
        return None

    @classmethod
    def getMachineName(cls):
        return cls.getMachineInfo().get("PRETTY_HOSTNAME", None)

    @classmethod
    def saveMachineInfo(cls, data):
        with open(cls.PATH, 'r') as f:
            lines = f.readlines()

        for key, value in data:
            assert key in ["CHASSIS", "PRETTY_HOSTNAME", "HARDWARE_VENDOR", "HARDWARE_MODEL"]

            bFound = False
            for i, line in enumerate(lines):
                if line.startswith(key + '='):
                    bFound = True
                    break

            if value is not None:
                if bFound:
                    lines[i] = key + '=' + value + '\n'
                else:
                    lines.append(key + '=' + value + '\n')
            else:
                if bFound:
                    lines.pop(i)

        with atomicwrites.atomic_write(cls.PATH, overwrite=True) as f:
            f.writelines(lines)

        cls._machineInfo = None

    @classmethod
    def saveChassisType(cls, chassis_type):
        _valueDict = {
            MachineType.DESKTOP: "desktop",
            MachineType.LAPTOP: "laptop",
            MachineType.TABLET: "tablet",
            MachineType.CONVERTIBLE: "convertible",
            MachineType.HANDSET: "handset",
            MachineType.WATCH: "watch",
            MachineType.SERVER: "server",
            MachineType.VM: "vm",
            MachineType.CONTAINER: "container",
        }

        with open(cls.PATH, 'r') as f:
            lines = f.readlines()

        bFound = False
        for i, line in enumerate(lines):
            if line.startswith('CHASSIS='):
                lines[i] = 'CHASSIS=' + _valueDict[chassis_type] + '\n'
                bFound = True
                break
        if not bFound:
            lines.append('CHASSIS=' + _valueDict[chassis_type] + '\n')

        with atomicwrites.atomic_write(cls.PATH, overwrite=True) as f:
            f.writelines(lines)

        cls._machineInfo = None

    @classmethod
    def saveMachineName(cls, machine_name):
        return cls.saveMachineInfo({
            "PRETTY_HOSTNAME", machine_name
        })


class MachineHwcfgFile:

    PATH = "/etc/machine-hwcfg"

    # sub objects
    _catDumb = None
    _catSaved = None
    _catProperties = None

    # json keys
    _keyDumb                = "dumb-components"
    _keySaved               = "saved-components"
    _keyMachineProperties   = "machine-properties"
    _keyComponentProperties = "component-properties"
    _keyDeviceProperties    = "device-properties"
    _keyConnectorProperties = "connector-properties"

    # cache for this file
    _machineHwCfgData = None

    @classmethod
    def catDumb(cls):
        if cls._catDumb is None:
            cls._catDumb = MachineHwcfgFileDumb(cls)
        return cls._catDumb

    @classmethod
    def catSaved(cls):
        if cls._catSaved is None:
            cls._catSaved = MachineHwcfgFileSaved(cls)
        return cls._catSaved

    @classmethod
    def catProperties(cls):
        if cls._catProperties is None:
            cls._catProperties = MachineHwcfgFileProperties(cls)
        return cls._catProperties

    @classmethod
    def save(cls):
        if cls._machineHwCfgData is None:
            cls._read()

        if len(cls._machineHwCfgData) == 0:
            Util.forceDelete(cls.PATH)
        else:
            with atomicwrites.atomic_write(cls.PATH, overwrite=True) as f:
                json.dump(cls._machineHwCfgData, f, indent=4)

    @classmethod
    def _read(cls):
        if cls._machineHwCfgData is None:
            if not os.path.exists(cls.PATH):
                cls._machineHwCfgData = {}
            else:
                with open(cls.PATH) as f:
                    data = json.load(f)
                    if True:
                        # check
                        data2 = data.copy()
                        data2.pop(cls._keyDumb, None)
                        data2.pop(cls._keySaved, None)
                        data2.pop(cls._keyMachineProperties, None)
                        data2.pop(cls._keyComponentProperties, None)
                        data2.pop(cls._keyDeviceProperties, None)
                        data2.pop(cls._keyConnectorProperties, None)
                        if len(data2) != 0:
                            raise ConfigError("Invalid content in %s" % (cls.PATH))
                    cls._machineHwCfgData = data

        return cls._machineHwCfgData


class MachineHwcfgFileDumb:

    """
    Example:
        {
            "dumb-components: [
                {
                    "vendor": "",
                    "model": "",
                    "serial": "",
                    "id": ""
                },
                {
                    "vendor": "",
                    "model": "",
                    "serial": ""
                },
                {
                    "vendor": "",
                    "model": "",
                    "id": ""
                },
                {
                    "vendor": "",
                    "model": ""
                }
            ],
        }
    """

    def __init__(self, parentClass):
        self._p = parentClass
        self._cache = None

    def read(self):
        if self._cache is None:
            data = self._p._read()
            if self._p._keyDumb in data:
                for c in data[self._p._keyDumb]:
                    if "vendor" not in c or "model" not in c:
                        raise ConfigError("Invalid dumb component in %s" % (self._p.PATH))

                    c2 = c.copy()
                    c2.pop("vendor")
                    c2.pop("model")
                    c2.pop("serial", None)
                    c2.pop("id", None)
                    if len(c2) > 0:
                        raise ConfigError("Invalid dumb component in %s" % (self._p.PATH))

                    if not all([isinstance(x, str) for x in c.values()]):
                        raise ConfigError("Invalid dumb component in %s" % (self._p.PATH))

                self._cache = data[self._p._keyDumb]

        return self._cache if self._cache is not None else []

    def add_dumb_component(self, vendor, model, serial=None, id=None):
        self._fillCache()

        c = {
            "vendor": vendor,
            "model": model,
        }
        if serial is not None:
            c["serial"] = serial
        if id is not None:
            c["id"] = id
        self._cache.append(c)

    def remove_dumb_component(self, **kwargs):
        assert self._cache is not None

        for i in reversed(range(len(self._cache))):
            c = self._cache[i]
            bMatch = True
            for k, v in kwargs.items():
                if k not in c:
                    bMatch = False
                    break
                if c[k] != v:
                    bMatch = False
                    break
            if bMatch:
                self._cache.pop(i)

        if len(self._cache) == 0:
            del self._p._read()[self._p._keyDumb]
            self._cache = None

    def _fillCache(self):
        if self._cache is None:
            data = self._p._read()
            data[self._p._keyDumb] = []
            self._cache = data[self._p._keyDumb]


class MachineHwcfgFileSaved:

    """
    Example:
        {
            "saved-components: [
                {
                    "vendor": "",
                    "model": "",
                    "save_info": {},
                    ...
                },
                {
                    "vendor": "",
                    "model": "",
                    "serial": "",
                    "save_info": {},
                    ...
                },
                {
                    "component_type": "xxx",
                    "save_info": {},
                    ...
                },
            ],
        }
    """

    def __init__(self, parentClass):
        self._p = parentClass

    def read(self):
        data = self._p._read()
        if self._p._keySaved in data:
            return data[self._p._keySaved]
        else:
            return []

    def flush(self, machine):
        data = self._p._read()

        if len(machine.get_components()) == 0:
            # same
            if self._p._keySaved not in data:
                return False

            # modified
            del data[self._p._keySaved]
            return True
        else:
            # generate savedComponentInfoData list
            savedComponentInfoDataList = []
            for c in machine.get_components():
                if c._getSaveInfo() is None:
                    continue

                assert "vendor" not in c._getSaveInfo()
                assert "model" not in c._getSaveInfo()
                assert "component_type" not in c._getSaveInfo()

                if c.get_vendor() is not None:
                    item = {
                        "vendor": c.get_vendor(),
                        "model": c.get_model(),
                    }
                else:
                    item = {
                        "component_type": c.get_type(),
                    }

                item.update(c._getSaveInfo())
                savedComponentInfoDataList.append(item)

            # same
            if self._p._keySaved in data:
                tlist = list(data[self._p._keySaved])
                try:
                    for savedComponentInfoData in savedComponentInfoDataList:
                        tlist.remove(savedComponentInfoData)
                    if len(tlist) == 0:
                        return False
                except ValueError:
                    pass

            # modified
            data[self._p._keySaved] = savedComponentInfoDataList
            return True


class MachineHwcfgFileProperties:

    """
    Example:
        {
            "machine-properties": [
                {
                    "apply": {
                        "serial": "xxx",
                    }
                }
            ],
            "component-properties": [
                {
                    "match": "xxx",
                    "apply": {
                        "serial": "xxx",
                        "default_offline_reason": "xxx",
                        "name": "xxx",
                    }
                },
                {
                    ...
                }
            ],
            "device-properties": [
                {
                    "match": "xxx",
                    "apply": {
                        "default_offline_reason": "xxx",
                        "name": "xxx"
                    }
                },
                {
                    ...
                }
            ],
            "connector-properties": [
                {
                    "match": "xxx",
                    "apply": {
                        "default_offline_reason": "xxx",
                        "peer_connector_matcher": "xxx",
                        "name": "xxx",
                    }
                },
                {
                    ...
                }
            ]
        }
        """

    def __init__(self, parentClass):
        self._p = parentClass

    def add_properties_for_machine_components_devices_connectors(self, machine):
        data = self._p._read()

        # machine
        v = data.get(self._p._keyMachineProperties, {}).get("apply", {}).get("serial", None)
        if v is not None:
            machine._setSerial(v)

        # components, devices and connectors
        if True:
            def __addPropertiesForComponent(apply, component):
                if "serial" in apply:
                    component._setSerial(apply["serial"])
                if "default_offline_reason" in apply:
                    if component.get_offline_reason() == OfflineReason.UNKNOWN:
                        component._setOfflineReason(OfflineReason(apply["default_offline_readon"]))
                if "name" in apply:
                    component._setName(apply["name"])

            def __addPropertiesForDevice(apply, device):
                if "default_offline_reason" in apply:
                    offlineReason = OfflineReason(apply["default_offline_reason"])
                    assert offlineReason != OfflineReason.REMOVED
                    if device.get_offline_reason() == OfflineReason.UNKNOWN:
                        device._setOfflineReason(offlineReason)
                if "name" in apply:
                    device._setName(apply["name"])

            def __addPropertiesForConnector(apply, connector):
                if "default_offline_reason" in apply:
                    offlineReason = OfflineReason(apply["default_offline_reason"])
                    assert offlineReason != OfflineReason.REMOVED
                    if connector.get_offline_reason() == OfflineReason.UNKNOWN:
                        connector._setOfflineReason(offlineReason)
                if "name" in apply:
                    connector._setName(apply["name"])

            mm = MatchMachine()
            for item in data.get(self._p._keyComponentProperties, []):
                mm.add_expression("component", item["mactcher"], functools.partial(__addPropertiesForComponent, item["apply"]))
            for item in data.get(self._p._keyDeviceProperties, []):
                mm.add_expression("device", item["mactcher"], functools.partial(__addPropertiesForDevice, item["apply"]))
            for item in data.get(self._p._keyConnectorProperties, []):
                mm.add_expression("connector", item["mactcher"], functools.partial(__addPropertiesForConnector, item["apply"]))
            mm.evaluate(machine)

    def machine_add_apply_serial(self, serial):
        assert serial is not None

        data = self._p._read()

        newValue = [
            {
                "apply": {
                    "serial": serial,
                }
            }
        ]
        if self._p._keyMachineProperties not in data or data[self._p._keyMachineProperties] != newValue:
            data[self._p._keyMachineProperties] = newValue
            return True
        else:
            return False

    def machine_remove_apply_serial(self):
        data = self._p._read()

        if self._p._keyMachineProperties in data:
            del data[self._p._keyMachineProperties]
            return True
        else:
            return False


class SeatUdevRule:

    @classmethod
    def generate(cls, seatName, components, devices):
        assert seatName == "seat-main"

        fullfn = cls._fullfn(seatName)

        buf = ""
        if True:
            buf += "# seat\n"
            buf += 'TAG=="seat", ENV{ID_SEAT}="%s"\n' % (seatName)
            buf += "\n"
            buf += "\n"

        return (fullfn, buf)

    @classmethod
    def write(cls, seatName, components, devices):
        fullfn, buf = cls.generate(seatName, components, devices)
        with open(fullfn, "w") as f:
            f.write(buf)

    @staticmethod
    def _fullfn(seatName):
        return "/etc/udev/rules.d/72-%s.rules" % (seatName)


class MatchMachine:

    def __init__(self):
        self._exprDict = {
            "component": dict(),
            "device": dict(),
            "connector": dict(),
        }

    def add_expression(self, category, expression, func):
        assert expression not in self._exprDict
        self._exprDict[category][expression] = func

    def evaluate(self, machine):
        self._machine = machine
        self._componentIndex = None
        self._index = None

        # components
        for expr, func in self._exprDict["component"].items():
            exprObj = Expression(expr)
            self._index = 0
            while self._index < self._machine._components:
                vars = [
                    ("component", self._machine._components[self._index]),
                ]
                if exprObj.evaluate(vars):
                    func(self._machine._components[self._index])
                self._index += 1

        # devices
        for expr, func in self._exprDict["device"].items():
            exprObj = Expression(expr)

            if len(self._machine._components) > 0:
                self._componentIndex = 0                # iterating component.device
            else:
                self._componentIndex = None             # iterating machine.device
            self._index = -1

            while self._index is not None:
                if self._componentIndex is not None:
                    vars = [
                        ("component", self._machine._components[self._componentIndex]),
                        ("device", self._machine._components[self._componentIndex]._getDevices()[self._index])
                    ]
                else:
                    vars = [
                        ("component", None),
                        ("device", self._machine._devices[self._index])
                    ]
                if exprObj.evaluate(vars):
                    if self._componentIndex is not None:
                        func(self._machine._components[self._componentIndex]._getDevices()[self._index])
                    else:
                        func(self._machine._devices[self._index])
                self._nextDevice()

        # connectors
        for expr, func in self._exprDict["connector"].items():
            exprObj = Expression(expr)

            if len(self._machine._components) > 0:
                self._componentIndex = 0                # iterating component.connector
            else:
                self._componentIndex = None             # iterating machine.connector
            self._index = -1

            while self._index is not None:
                if self._componentIndex is not None:
                    vars = [
                        ("component", self._machine._components[self._componentIndex]),
                        ("connector", self._machine._components[self._componentIndex]._getConnectors()[self._index])
                    ]
                else:
                    vars = [
                        ("component", None),
                        ("connector", self._machine._connectors[self._index])
                    ]
                if exprObj.evaluate(vars):
                    if self._componentIndex is not None:
                        func(self._machine._components[self._componentIndex]._getConnectors()[self._index])
                    else:
                        func(self._machine._connectors[self._index])
                self._next()

        del self._componentIndex
        del self._index
        del self._machine

    def _nextDevice(self):
        while True:
            if self._componentIndex is not None:
                if self._index < len(self._machine._components[self._componentIndex]._getDevices()) - 1:
                    self._index += 1
                    break
                else:
                    if self._componentIndex < len(self._machine._components) - 1:
                        self._componentIndex += 1
                        self._index = -1
                        continue
                    else:
                        self._componentIndex = None
                        self._index = -1
                        continue
            else:
                if self._index < len(self._machine._devices) - 1:
                    self._index += 1
                    break
                else:
                    self._index = None
                    break

    def _nextConnector(self):
        while True:
            if self._componentIndex is not None:
                if self._index < len(self._machine._components[self._componentIndex]._getDevices()) - 1:
                    self._index += 1
                    break
                else:
                    if self._componentIndex < len(self._machine._components) - 1:
                        self._componentIndex += 1
                        self._index = -1
                        continue
                    else:
                        self._componentIndex = None
                        self._index = -1
                        continue
            else:
                if self._index < len(self._machine._connectors) - 1:
                    self._index += 1
                    break
                else:
                    self._index = None
                    break


class MatchComponent:

    def __init__(self):
        self._exprDict = dict()

    def add_expression(self, expression, func):
        assert expression not in self._exprDict
        self._exprDict[expression] = func

    def evaluate(self, component):
        bOk = False
        for expr, func in self._exprDict.items():
            if Expression(expr).evaluate([("component", component)]):
                assert not bOk
                func(self._component)
                bOk = True


class MatchDevice:

    def __init__(self):
        self._exprDict = dict()

    def add_expression(self, expression, func):
        assert expression not in self._exprDict
        self._exprDict[expression] = func

    def evaluate(self, device):
        bOk = False
        for expr, func in self._exprDict.items():
            if Expression(expr).evaluate([("device", device)]):
                assert not bOk
                func(device)
                bOk = True


class MatchConnector:

    def __init__(self):
        self._exprDict = dict()

    def add_expression(self, expression, func):
        assert expression not in self._exprDict
        self._exprDict[expression] = func

    def evaluate(self, connector):
        bOk = False
        for expr, func in self._exprDict.items():
            if Expression(expr).evaluate([("connector", connector)]):
                assert not bOk
                func(self._connector)
                bOk = True
